import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import cv2 as cv
import numpy as np
from skimage.feature import local_binary_pattern, hog

class FeatureExtractor:
    def __init__(self, root):
        self.root = root
        self.root.title("图形特征提取工具")
        self.root.geometry("1000x700")
        
        # 初始化变量
        self.original_img = None  # 保存原始图像
        self.current_img = None
        self.display_img = None
        self.photo = None
        
        # 特征提取算法选项
        self.algorithms = {
            "边缘检测 (Canny)": self.extract_edges,
            "角点检测 (Harris)": self.extract_corners,
            "Shi-Tomasi 角点检测": self.extract_shitomasi,
            "SIFT 特征": self.extract_sift,
            "SURF 特征": self.extract_surf,
            "FAST 角点检测": self.extract_fast,
            "ORB 特征检测": self.extract_orb,
            "LBP 原始模式": lambda: self.extract_lbp(mode='default'),
            "LBP 圆形模式": lambda: self.extract_lbp(mode='ror'),
            "LBP 旋转不变模式": lambda: self.extract_lbp(mode='ror'),
            "LBP 特价模式": lambda: self.extract_lbp(mode='uniform'),
            "HOG 特征提取": self.extract_hog
        }
        
        self.setup_ui()
    
    def setup_ui(self):
        # 创建菜单栏
        menubar = tk.Menu(self.root)
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="打开图像", command=self.open_image)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.root.quit)
        menubar.add_cascade(label="文件", menu=file_menu)
        self.root.config(menu=menubar)
        
        # 创建主框架
        main_frame = tk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 创建顶部框架（图像显示区域）
        top_frame = tk.Frame(main_frame)
        top_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建画布
        self.canvas = tk.Canvas(top_frame, bg='lightgray')
        self.canvas.pack(fill=tk.BOTH, expand=True)
        
        # 创建底部控制面板
        bottom_frame = tk.Frame(main_frame)
        bottom_frame.pack(fill=tk.X, pady=10)
        
        # 算法选择下拉菜单
        tk.Label(bottom_frame, text="选择特征提取算法:").pack(side=tk.LEFT, padx=5)
        self.algorithm_var = tk.StringVar()
        self.algorithm_var.set("边缘检测 (Canny)")  # 默认算法
        algorithm_menu = tk.OptionMenu(bottom_frame, self.algorithm_var, *self.algorithms.keys())
        algorithm_menu.pack(side=tk.LEFT, padx=5)
        
        # 提取按钮
        extract_button = tk.Button(bottom_frame, text="提取特征", command=self.extract_features)
        extract_button.pack(side=tk.LEFT, padx=5)
        
        # 还原按钮
        reset_button = tk.Button(bottom_frame, text="还原", command=self.reset_image)
        reset_button.pack(side=tk.LEFT, padx=5)
        
        # 状态栏
        self.status_bar = tk.Label(self.root, text="就绪", bd=1, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
    
    def open_image(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp *.gif")]
        )
        
        if file_path:
            # 读取图像
            img = cv.imread(file_path)
            if img is not None:
                self.original_img = img.copy()  # 保存原始图像
                self.current_img = img
                self.display_image_on_canvas()
                self.status_bar.config(text=f"已打开图像: {file_path}")
            else:
                tk.messagebox.showerror("错误", "无法读取图像文件")
    
    def display_image_on_canvas(self):
        if self.current_img is None:
            return
            
        # 获取画布尺寸
        canvas_width = self.canvas.winfo_width()
        canvas_height = self.canvas.winfo_height()
        
        # 如果画布尺寸为0，使用默认值
        if canvas_width <= 1 or canvas_height <= 1:
            canvas_width, canvas_height = 800, 500
            
        # 获取图像尺寸
        img_h, img_w = self.current_img.shape[:2]
        
        # 计算缩放比例
        scale = min(canvas_width/img_w, canvas_height/img_h) * 0.9  # 留一些边距
        display_w = int(img_w * scale)
        display_h = int(img_h * scale)
        
        # 缩放图像
        self.display_img = cv.resize(self.current_img, (display_w, display_h))
        
        # 转换为RGB格式
        if len(self.display_img.shape) == 3:
            rgb_img = cv.cvtColor(self.display_img, cv.COLOR_BGR2RGB)
        else:
            rgb_img = cv.cvtColor(self.display_img, cv.COLOR_GRAY2RGB)
            
        # 转换为PIL图像并在画布上显示
        pil_img = Image.fromarray(rgb_img)
        self.photo = ImageTk.PhotoImage(pil_img)
        
        # 清空画布
        self.canvas.delete("all")
        
        # 在画布中心显示图像
        self.canvas.create_image(
            (canvas_width - display_w) // 2,
            (canvas_height - display_h) // 2,
            anchor=tk.NW, image=self.photo
        )
    
    def extract_features(self):
        if self.current_img is None:
            tk.messagebox.showwarning("警告", "请先打开图像")
            return
        
        # 获取选择的算法
        selected_algorithm = self.algorithm_var.get()
        algorithm_func = self.algorithms.get(selected_algorithm)
        
        if algorithm_func:
            algorithm_func()
    
    def extract_edges(self):
        """使用 Canny 算法提取边缘"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        edges = cv.Canny(gray_img, 100, 200)
        self.display_result(edges, "边缘检测结果")
    
    def extract_corners(self):
        """使用 Harris 算法提取角点"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        gray_img = np.float32(gray_img)
        corners = cv.cornerHarris(gray_img, 2, 3, 0.04)
        corners = cv.dilate(corners, None)
        self.current_img[corners > 0.01 * corners.max()] = [0, 0, 255]  # 标记角点为红色
        self.display_image_on_canvas()
        self.status_bar.config(text="角点检测完成 (红色标记)")
    
    def extract_sift(self):
        """提取 SIFT 特征"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        sift = cv.SIFT_create()
        keypoints, descriptors = sift.detectAndCompute(gray_img, None)
        result_img = cv.drawKeypoints(self.current_img, keypoints, None, color=(0, 255, 0))
        self.display_result(result_img, "SIFT 特征提取结果")
    
    def extract_surf(self):
        """提取 SIFT 特征（替代 SURF）"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        sift = cv.SIFT_create()
        keypoints, descriptors = sift.detectAndCompute(gray_img, None)
        result_img = cv.drawKeypoints(self.current_img, keypoints, None, color=(0, 255, 0))
        self.display_result(result_img, "SIFT 特征提取结果（替代 SURF）")
    
    def extract_shitomasi(self):
        """使用 Shi-Tomasi 算法提取角点"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        corners = cv.goodFeaturesToTrack(gray_img, maxCorners=100, qualityLevel=0.01, minDistance=10)
        corners = np.int0(corners)
        result_img = self.current_img.copy()
        for corner in corners:
            x, y = corner.ravel()
            cv.circle(result_img, (x, y), 3, (0, 255, 0), -1)  # 标记角点为绿色
        self.display_result(result_img, "Shi-Tomasi 角点检测结果")
    
    def extract_fast(self):
        """使用 FAST 算法提取角点"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        fast = cv.FastFeatureDetector_create()
        keypoints = fast.detect(gray_img, None)
        
    def extract_hog(self):
        """提取 HOG 特征"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        hog_features, hog_image = hog(gray_img, orientations=8, pixels_per_cell=(16, 16),
                                    cells_per_block=(1, 1), visualize=True)
        hog_image = cv.normalize(hog_image, None, 0, 255, cv.NORM_MINMAX, dtype=cv.CV_8U)
        self.display_result(hog_image, "HOG 特征提取结果")
        result_img = cv.drawKeypoints(self.current_img, keypoints, None, color=(0, 255, 0))
        self.display_result(result_img, "FAST 角点检测结果")
    
    def extract_orb(self):
        """使用 ORB 算法提取特征"""
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        orb = cv.ORB_create()
        keypoints, descriptors = orb.detectAndCompute(gray_img, None)
        result_img = cv.drawKeypoints(self.current_img, keypoints, None, color=(0, 255, 0))
        self.display_result(result_img, "ORB 特征检测结果")
    
    def extract_lbp(self, mode='default'):
        """使用 LBP 算法提取特征
        :param mode: LBP 模式 ('default', 'circular', 'ror', 'uniform')
        """
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        
        if mode == 'default':
            lbp = local_binary_pattern(gray_img, P=8, R=1, method='default')
            title = "LBP 原始模式"
        elif mode == 'ror':
            lbp = local_binary_pattern(gray_img, P=8, R=1, method='ror')
            title = "LBP 旋转不变模式"
        elif mode == 'ror':
            lbp = local_binary_pattern(gray_img, P=8, R=1, method='ror')
            title = "LBP 旋转不变模式"
        elif mode == 'uniform':
            lbp = local_binary_pattern(gray_img, P=8, R=1, method='uniform')
            title = "LBP 特价模式"
        else:
            lbp = local_binary_pattern(gray_img, P=8, R=1, method='default')
            title = "LBP 原始模式"
        
        result_img = cv.normalize(lbp, None, 0, 255, cv.NORM_MINMAX, dtype=cv.CV_8U)
        self.display_result(result_img, title)
        gray_img = cv.cvtColor(self.current_img, cv.COLOR_BGR2GRAY)
        orb = cv.ORB_create()
        keypoints, descriptors = orb.detectAndCompute(gray_img, None)
        result_img = cv.drawKeypoints(self.current_img, keypoints, None, color=(0, 255, 0))
        self.display_result(result_img, "ORB 特征检测结果")
    
    def display_result(self, result_img, title):
        """显示特征提取结果"""
        # 转换为RGB格式
        if len(result_img.shape) == 3:
            rgb_img = cv.cvtColor(result_img, cv.COLOR_BGR2RGB)
        else:
            rgb_img = cv.cvtColor(result_img, cv.COLOR_GRAY2RGB)
        
        # 转换为PIL图像并显示
        pil_img = Image.fromarray(rgb_img)
        photo = ImageTk.PhotoImage(pil_img)
        
        # 清空画布
        self.canvas.delete("all")
        
        # 在画布中心显示结果
        self.canvas.create_image(
            (self.canvas.winfo_width() - result_img.shape[1]) // 2,
            (self.canvas.winfo_height() - result_img.shape[0]) // 2,
            anchor=tk.NW, image=photo
        )
        
        # 保持引用
        self.photo = photo
        self.status_bar.config(text=title)
    
    def reset_image(self):
        """还原图像到初始状态"""
        if self.original_img is not None:
            self.current_img = self.original_img.copy()
            self.display_image_on_canvas()
            self.status_bar.config(text="图像已还原")


if __name__ == "__main__":
    root = tk.Tk()
    app = FeatureExtractor(root)
    root.mainloop()